package lc_721_账户合并;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeSet;

public class Solution {

    class UnionFind {

        int[] p;

        int stack[], sp, sizes[];

        public UnionFind(int size) {
            p = new int[size];
            stack = new int[size];
            sizes = new int[size];
            for (int i = 0; i < p.length; i++) {
                p[i] = i;
                sizes[i] = 1;
            }
        }

        public void union(int a, int b) {
            int pa = find(a), pb = find(b);
            if (pa != pb) {
                if (sizes[pb] < sizes[pa]) {
                    sizes[pa] += sizes[pb];
                    sizes[pb] = 0;
                    p[pb] = pa;
                } else {
                    sizes[pb] += sizes[pa];
                    sizes[pa] = 0;
                    p[pa] = pb;
                }
            }
        }

        public int find(int id) {
            sp = 0;
            while (p[id] != id) {
                stack[sp++] = id;
                id = p[id];
            }
            while (sp > 0) {
                p[stack[--sp]] = id;
            }
            return p[id];
        }

        public boolean isSameSet(int a, int b) {
            return find(a) == find(b);
        }
    }

    public List<List<String>> accountsMerge(List<List<String>> accounts) {

        UnionFind uf = new UnionFind(accounts.size());
        Map<String, Integer> email_p = new HashMap<>();
        for (int i = 0; i < accounts.size(); i++) {
            List<String> account = accounts.get(i);
            for (int j = 1; j < account.size(); j++) {
                int id = email_p.getOrDefault(account.get(j), -1);
                if (id == -1) {
                    email_p.put(account.get(j), i);
                } else {
                    uf.union(i, id);
                }
            }
        }
        Map<Integer, TreeSet<String>> id_emails = new HashMap<>();
        for (int i = 0; i < accounts.size(); i++) {
            List<String> account = accounts.get(i);
            for (int j = 1; j < account.size(); j++) {
                id_emails.computeIfAbsent(uf.find(email_p.get(account.get(j))), k -> new TreeSet<>())
                        .add(account.get(j));
            }
        }
        List<List<String>> ans = new ArrayList<>();
        for (Map.Entry<Integer, TreeSet<String>> entry : id_emails.entrySet()) {
            List<String> account = new ArrayList<>();
            account.add(accounts.get(entry.getKey()).get(0));
            account.addAll(entry.getValue());
            ans.add(account);
        }
        return ans;
    }
}